import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.svm import SVC
import numpy as np


class CNNSVM(nn.Module):
    def __init__(self, device=None, n_classes=2):
        super().__init__()
        self.device = device or torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.n_classes = n_classes
        
        # 使用更灵活的特征提取结构，能够适应不同的输入尺寸
        # 第一阶段：时间维度卷积
        self.time_conv = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=(1, 25), padding=(0, 12)),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.MaxPool2d((1, 4), (1, 4))
        )
        
        # 空间卷积将在forward中动态创建
        self.spatial_conv = None
        
        # 特征处理层
        self.feature_processing = nn.Sequential(
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d((1, 4), (1, 4)),
            nn.Dropout(0.25),
            nn.Conv2d(64, 128, kernel_size=(1, 10), padding=(0, 5)),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d(1)  # 全局平均池化到固定大小
        )
        
        # SVM配置
        self.clf = make_pipeline(
            StandardScaler(), 
            SVC(
                C=10,
                kernel='rbf',
                gamma='scale',
                probability=True,
                class_weight='balanced',
                decision_function_shape='ovr'
            )
        )
        
        self.is_fitted = False
        self.features_shape = None

    def forward(self, x, return_features=False, apply_activation=True):
        try:
            batch_size, _, num_electrodes, time_samples = x.shape
            
            # 第一阶段：时间维度卷积
            x = self.time_conv(x)  # [batch, 32, num_electrodes, time/4]
            
            # 动态创建或更新空间卷积层
            if self.spatial_conv is None or self.spatial_conv[0].weight.size(2) != num_electrodes:
                self.spatial_conv = nn.Sequential(
                    nn.Conv2d(32, 64, kernel_size=(num_electrodes, 1), stride=(1, 1)),
                    nn.BatchNorm2d(64),
                    nn.ReLU(),
                    nn.Dropout(0.25)
                ).to(self.device)
            
            # 应用空间卷积
            x = self.spatial_conv(x)  # [batch, 64, 1, time/4]
            
            # 应用特征处理
            features = self.feature_processing(x)  # [batch, 128, 1, 1]
            features = features.view(batch_size, -1)  # [batch, 128]
            
            if self.features_shape is None:
                self.features_shape = features.shape[1]
                
            if return_features:
                return features
                
            # 训练模式或SVM未拟合，返回特征或占位符
            if self.training or not self.is_fitted:
                # 如果在训练过程中或SVM尚未拟合，我们只返回特征或占位符
                if self.n_classes == 2:
                    # 为二分类返回单个值
                    zeros = torch.zeros(batch_size, 1, device=self.device)
                    return zeros  # 返回占位符值
                else:
                    # 为多分类返回零数组
                    zeros = torch.zeros(batch_size, self.n_classes, device=self.device)
                    return zeros
            
            # 推理模式，使用已拟合的SVM
            features_np = features.detach().cpu().numpy()
            
            # 获取SVM预测
            if self.n_classes == 2:
                # 二分类
                try:
                    if not apply_activation:
                        # 返回未激活的分数（尽管SVM预测已经是概率）
                        predictions = self.clf.decision_function(features_np)
                        return torch.tensor(predictions, device=self.device).unsqueeze(1)
                    else:
                        # 返回概率
                        predictions = self.clf.predict_proba(features_np)[:, 1]
                        return torch.tensor(predictions, device=self.device).unsqueeze(1)
                except:
                    # 如果预测失败，返回默认值
                    if not apply_activation:
                        return torch.zeros(batch_size, 1, device=self.device)
                    else:
                        return torch.zeros(batch_size, 1, device=self.device)
            else:
                # 多分类
                try:
                    if not apply_activation:
                        # 返回未激活的分数（尽管SVM预测已经是概率）
                        # 对于多分类，我们需要使用decision_function或predict_proba
                        try:
                            scores = self.clf.decision_function(features_np)
                        except:
                            # 如果不支持decision_function，使用predict_proba并转换
                            probs = self.clf.predict_proba(features_np)
                            # 概率转换为分数的简单方法（不是真正的logits，但可以作为近似）
                            scores = np.log(probs + 1e-10)
                        return torch.tensor(scores, device=self.device)
                    else:
                        # 返回概率
                        predictions = self.clf.predict_proba(features_np)
                        return torch.tensor(predictions, device=self.device)
                except:
                    # 如果预测失败，返回均匀分布或零
                    if not apply_activation:
                        return torch.zeros(batch_size, self.n_classes, device=self.device)
                    else:
                        return torch.ones(batch_size, self.n_classes, device=self.device) / self.n_classes
        
        except Exception as e:
            print(f"CNNSVM forward传播错误: {str(e)}")
            print(f"输入x形状: {x.shape}")
            
            # 返回默认值
            if return_features:
                return torch.zeros(batch_size, 128, device=self.device)
            elif not apply_activation:
                if self.n_classes == 2:
                    return torch.zeros(batch_size, 1, device=self.device)
                else:
                    return torch.zeros(batch_size, self.n_classes, device=self.device)
            elif self.n_classes == 2:
                return torch.zeros(batch_size, 1, device=self.device)
            else:
                return torch.ones(batch_size, self.n_classes, device=self.device) / self.n_classes
            
    def fit_svm(self, features, labels):
        """
        Fit the SVM classifier with extracted features
        
        Args:
            features: Tensor of shape (n_samples, feature_dim)
            labels: Tensor of shape (n_samples,)
        """
        try:
            # Convert to numpy
            features_np = features.detach().cpu().numpy()
            labels_np = labels.detach().cpu().numpy()
            
            # Fit SVM
            self.clf.fit(features_np, labels_np)
            self.is_fitted = True
            
            return self
        except Exception as e:
            print(f"SVM拟合错误: {str(e)}")
            return self

